import asyncio
import os
import datetime
import json
from datetime import date, timedelta
from io import StringIO

import numpy as np
import pandas as pd
import streamlit as st

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import exc

from db.forecast import get_forecast_by_id
from models.models import Schedule, ScheduleName
from tasks.task_manager import task_queue


def read_prepare_doctors(file_name):
    doctors = pd.read_csv(file_name, sep = ';').fillna('')
    doctors.columns = ['name', 'fte', 'modality','secondary']
    mod_list = ['Денс', 'КТ', 'ММГ', 'МРТ', 'РГ', 'Флг']
    doctors['sec_list'] = doctors.secondary.apply(lambda x: [z for z in [ y.strip() for y in x.split(',')] if len(z)>0 ] )
    for c in mod_list:
        doctors[c] = doctors.sec_list.apply(lambda x: 1 if c in x else 0)
    doctors['num_mod'] = doctors.sec_list.apply(lambda x: len(x)+1)
    return doctors.drop(columns=['sec_list']).sort_values(['num_mod','name'],ascending=False).reset_index(drop=True)

#
# Блоки фиксированного расписания
#
SHED_BLOCK = [1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,
            0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,
            0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,
            1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,
            1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,
            1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,
            1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1,1,1,1,1,0,0,1]

SHED_SL = [0,0,1,0,0,0,0,0,0,1,0,0,0,0,0,0,1,0,0,0,0,0,0,1,0,0,0,0,0,0,1,0,0,0,0,
           0,0,1,1,0,0,0,0,0,1,1,0,0,0,0,0,1,1,0,0,0,0,0,1,1,0,0,0,0,0,1,1,0,0,0,
           0,1,1,1,0,0,0,0,1,1,1,0,0,0,0,1,1,1,0,0,0,0,1,1,1,0,0,0,0,1,1,1,0,0,0,
           0,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,
           1,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,0,1,1,1,1,0,0,]

shed_block = np.array(SHED_BLOCK).reshape((7,-1))
shed_line = np.array(SHED_SL).reshape((5,-1))
#
#   Составление расписания по количеству часов
#
def doctors_in_need(workhours):
    #print(workhours)
    num_fixed = workhours // 40 
    variable = 1 if workhours % 40 >0 else 0
    return num_fixed, variable
    
def plan_block_copy(num):
    nc = num #// 40
    l, c = [], 0
    for i in range(nc):
        l.append((i,c))
        c = c+1 if c < 6 else 0
    return l
    
def make_shedule(workhours):
    workhours = np.round(workhours/8.0).astype(np.int32)*8
    #print(workhours)
    num_fixed, variable = doctors_in_need(workhours) 
    var_ind = (workhours - num_fixed * 40) // 8 - 1
    shedule = np.zeros((num_fixed+variable,35))
    if num_fixed > 0: 
        for i, j in plan_block_copy(num_fixed):
            shedule[i,:] = shed_block[j,:]
    if variable >0 :
        shedule[num_fixed, : ] = shed_line[var_ind,:]
    return shedule
#
# шаблон расписания
#
def prepare_template(data):
    num_days = len(data.week.unique())*7
    start_date = data.date.min()
    days = np.arange(num_days)
    t = pd.DataFrame(days, columns = ['day'])
    t['date'] = t.day.apply(lambda x: start_date+timedelta(days=x))
    t['year'] = t.date.dt.year
    t['week'] = t.date.apply(lambda x: x.isocalendar()[1])
    return t.drop(columns=['day'])
#
#
#
def notify(to_std,text):
    if to_std : print(text)
    else : st.write(text)
#
#  Подготовка расписания по исходному файлу data
#  Файл data должен иметь размер по необходимому числу недель, например 5
#
#
def prepare_shedule(data, doctor_file, to_std, task_name: str):
    doctors = read_prepare_doctors(doctor_file)
    
    # Формирование задачи на составление расписания
    task = (data.groupby(['week','modality'], as_index=False).agg({'workhours':'sum'})
        .groupby(['modality'], as_index=False).agg({'workhours':'mean'})
       .sort_values('workhours').reset_index(drop=True))
    task['workhours'] = np.round(task.workhours).astype(np.int32)
    
    # Составление расписания
    utilised_doctors = []
    templ = prepare_template(data)
    shed_list = []
    progress = 10
    for i in range(len(task)):
        modality = task.at[i,'modality']; workhours = task.at[i,'workhours']
        shedule = make_shedule(workhours)
        num_req_doctors = shedule.shape[0]
        remain_doctors = doctors[~doctors.index.isin(utilised_doctors)]
        avaliable = list(remain_doctors[remain_doctors.modality==modality].index) + list(remain_doctors[remain_doctors[modality] == 1].index)
        selected = avaliable[:min(num_req_doctors,len(avaliable))]
        wl = []
        for j in range(len(selected)): 
            t = templ.copy()
            t['name'] = doctors.at[selected[j],'name']
            t['modality'] = modality
            t['workhours'] = (shedule[j,:] * 8).astype(np.int16)
            t.drop( index = t[t.workhours == 0].index, inplace = True)
            wl.append(t)
        shed_list.append(pd.concat(wl, axis=0).sort_values(by=['date','name']).reset_index(drop=True))
        utilised_doctors += selected
        # notify(to_std, f" **{modality}**: Прогнозируется {workhours} часов. Требуется {num_req_doctors} специалистов, доступно {len(avaliable)}, использовано {len(selected)}." )
        task_queue.put(
            item={
                "task_type": "schedule",
                "task_name": task_name,
                "progress": progress,
                "msg": f" **{modality}**: Прогнозируется {workhours} часов. Требуется {num_req_doctors} специалистов, доступно {len(avaliable)}, использовано {len(selected)}.",
            }
        )
        progress += i
    # notify(to_std,f"Осталось неиспользовано {len(doctors)-len(utilised_doctors)} специалистов. ")
    task_queue.put(
            item={
                "task_type": "schedule",
                "task_name": task_name,
                "progress": progress,
                "msg": f"Осталось неиспользовано {len(doctors)-len(utilised_doctors)} специалистов.",
            }
        )
    return pd.concat(shed_list, axis=0).reset_index(drop=True)  #sort_values(by=['date', 'modality', 'name']).


async def task_create_schedule(schedule_name: str, forecast_id: int, schedule_start):
    """Задача формирования расписания."""
    SQLALCHEMY_DATABASE_URL = "postgresql+asyncpg://user:pass@postgres_db:5432/dbname"
    engine = create_async_engine(
        SQLALCHEMY_DATABASE_URL,
    )

    async_session = sessionmaker(
        engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )

    task_queue.put(
        item={
            "task_type": "schedule",
            "task_name": schedule_name,
            "status": "DOING",
            "progress": 5,
            "msg": "Начало загрузки данных...",
        }
    )
    async with async_session() as session:
        all_forecast = await get_forecast_by_id(
            forecast_id=forecast_id, session=session
        )

    async with async_session() as session:
        new_schedule_name = ScheduleName(name=schedule_name)
        session.add(new_schedule_name)
        try:
            await session.commit()
        except exc.IntegrityError:
            await session.rollback()
            task_queue.put(
                item={
                    "task_type": "schedule",
                    "task_name": schedule_name,
                    "status": "ERROR",
                    "progress": 0,
                    "error": "Расписание уже существует в базе данных",
                    
                }
            )
            return

    for item in all_forecast:
        if isinstance(item["date"], date):
            item["date"] = datetime.datetime.strftime(
                item["date"], "%Y-%m-%d"
            )  # item["date"]

    all_forecast_json = StringIO(json.dumps(all_forecast))

    forecast = pd.read_json(all_forecast_json, orient="records", convert_dates=False)
    forecast_week_starts = sorted(list(forecast.date.unique()))

    NUM_WEEKS = 5
    try:
        schedule_end = forecast_week_starts[forecast_week_starts.index(schedule_start)+NUM_WEEKS-1]
    except (IndexError, ValueError):
        task_queue.put(
                item={
                    "task_type": "schedule",
                    "task_name": schedule_name,
                    "status": "ERROR",
                    "progress": 0,
                    "error": "Ошибка в дате начала расписания.",
                    
                }
            )
        return
    selection = forecast[(forecast.date>=schedule_start)&(forecast.date<=schedule_end)]
    selection['date'] = pd.to_datetime(selection['date'])

    doctor_file = os.path.join("download", "doctors_list.txt")

    shedule = prepare_shedule(selection, doctor_file, to_std=False, task_name=schedule_name)
    
    # shedule.to_csv(schedule_name, index = False)
    
    records = shedule.to_dict(orient="records")
    # st.write('Расписание записано')
    weekday = {
        1: "Понедельник",
        2: "Вторник",
        3: "Среда",
        4: "Четверг",
        5: "Пятница",
        6: "Суббота",
        7: "Воскресенье",
    }
    async with async_session() as session:
        for record in records:
            # print('record', record)
            schedule = Schedule(
                schedule_name=schedule_name,
                date=record["date"].date(),
                name=record["name"],
                modality=record["modality"],
                week=int(record["week"]),
                # weekday=weekday.get(record["weekday"]),
                workhours=record["workhours"],
            )
            session.add(schedule)
        await session.commit()
    task_queue.put(
        item={
            "task_type": "schedule",
            "task_name": schedule_name,
            "progress": 100,
            "msg": "Расписание записано",
            "status": "DONE",
        }
    )




def start_schedule_task(
    schedule_name: str,
    forecast_id: int,
    date_start_schedule: date,
):
    """Запускает асинхронную функцию в отдельном потоке."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(
        task_create_schedule(
            schedule_name=schedule_name,
            forecast_id=forecast_id,
            schedule_start=date_start_schedule,
        )
    )
